package com.arron.algorithm.leetcodetop100.二叉树;

import com.arron.algorithm.tree.TreeNode;

import java.util.HashMap;
import java.util.Map;

/**
 *  230. 二叉搜索树中第K小的元素
 */
public class kthSmallest {

    int count = 0;

    int t = 0;
    int res = 0;
    TreeNode treeNode = null;

    Map<TreeNode,Integer> rootNodeCountMap = new HashMap<>();

    /**
     *  递归
     * @param root
     * @param k
     * @return
     */
    public int kthSmallest(TreeNode root, int k) {

        if (root == null) {
            return 0;
        }
        t = k;
       kthSmallest(root);
        return res;
    }


    public void kthSmallest(TreeNode root){
        if (root == null) return;
        kthSmallest(root.left);
        count++;
        if (count == t){
            res = root.val;
            return;
        }
        kthSmallest(root.right);
    }


    /**
     *  优化版，利用二叉搜索树的特性 ，每个节点记录包括自己的节点个数
     */
    public TreeNode kthSmallest_plus(TreeNode root,int k){

        if (root == null){
            return null;
        }
        TreeNode temp =root;

        TreeNode res = null;
        while (temp !=null){

            int leftNums = getNodeNums(temp.left);

            if (leftNums +1 == k){
                res = temp;
                break;
            }else if (leftNums+1 <k){
                k = k-leftNums+1;
                temp = temp.right;
            }else {
                temp = temp.left;
            }
        }
        return res;
    }


    //统计每个节点的 以自己为根节点的子节点个数（包括自己）,放到hashmap中方便访问
    public int countNodeNums(TreeNode root){
        if (root == null) {
            return 0;
        }
//        int leftCount =  countNodeNums(root.left);
//        int rightCount = countNodeNums(root.right);
//
//        rootNodeCountMap.put(root,leftCount+rightCount+1);

        //优化一下
        rootNodeCountMap.put(root, countNodeNums(root.left)+ countNodeNums(root.right)+1);
        return rootNodeCountMap.get(root);
//        return leftCount+rightCount+1;
    }

    public int getNodeNums(TreeNode root){
        return rootNodeCountMap.getOrDefault(root,0);
    }


    public static void main(String[] args) {

        //测试 countNodeNums 方法
        BuildTree buildTree = new BuildTree();
        TreeNode treeNode = buildTree.buildTree();
        kthSmallest kthSmallest = new kthSmallest();
        kthSmallest.countNodeNums(treeNode);

        System.out.println(kthSmallest.rootNodeCountMap);


    }



}
